/*!
\copyright  Copyright (c) 2017 Qualcomm Technologies International, Ltd.
            All Rights Reserved.
            Qualcomm Technologies International, Ltd. Confidential and Proprietary.
\version
\file       av_headset_upgrade.c
\brief      Application upgrade management.

Over the air upgrade is managed from this file. Interaction
with the headset software is required as well as support from
the \b upgrade and \b gaia VM libraries.

This is a minimal implementation of upgrade.
*/

#ifdef INCLUDE_USB_HID_UPGRADE

#include "source_upgrade.h"
#include "source_usb.h"
#include "source_ps.h"
#include "source_debug.h"
#include "secrsa_padding.h"
#include <vmal.h>
#include <panic.h>
#include <string.h>
#include <stdlib.h>
#include <boot.h>
#include <ctype.h>
#include <stdio.h>
#include <source.h>
#include <stream.h>
#include <partition.h>
#include <imageupgrade.h>
#include <logging.h>
#include <vmal.h>
#include <operator.h>
#include <rsa_pss_constants.h>
#include <hydra_macros.h>
#include <rsa_decrypt.h>

#define ACK_MASK 0x80
#define NACK_MASK 0
#define ACK TRUE
#define NACK FALSE
#define MAX_IMAGE_DATA_PAYLOAD    59
#define PARTITION_SIGNATURE (IMAGE_SECTION_ID_MAX + 1)
/* The number of uint8s in a SHA-256 signature */
#define SHA_256_HASH_LENGTH (256/8)
#define RSA_SIGNATURE_SIZE (128) /* 128 * 16 = 2048 bits */

enum
{
    UPGRADE_ERASE_FLASH	= 1,
    UPGRADE_SELECT_PARTITION,
    UPGRADE_WRITE_IMAGE,
    UPGRADE_VALIDATE_IMAGE,
    UPGRADE_READ_VERSION,
    UPGRADE_WRITE_VERSION,
    UPGRADE_COPY_MISSING_SECTIONS,
};

typedef enum
{
    UPGRADE_STATE_IDLE	= 0,
    UPGRADE_STATE_START,
    UPGRADE_STATE_ERASING_FLASH,
    UPGRADE_STATE_DOWNLOADING,
    UPGRADE_STATE_IMAGE_VALIDATED
} UPGRADE_STATE_T;

typedef struct
{
    uint8 mainCmd;
    uint8 attr;
    uint8 data[61];
} UPGRADE_CMD_T;

typedef struct {
    uint8 partitionNumber;
    uint8 partitionSize[4];
    uint8 checkSum[4];
    uint8 firstWord[4];
} UpgradePartitionSelectMsg_t;

typedef struct {
    uint8 dataSize;
    uint8 lastPacket;
    uint8 data[MAX_IMAGE_DATA_PAYLOAD];
} UpgradeImageWriteMsg_t;

typedef struct
{
    TaskData  upgradeTask;
    Sink upgradeSink;
    uint8 selectedPartition;
    uint32 selectedpartitionCheckSum;
    uint32 selectedpartitionSize;
    uint32 currentPartitionCheckSum;
    uint32 currentPartitionPos;
    uint8 *signature;
    uint32 signatureSize;
    uint32 partitionWritten;
    hash_context_t vctx;
    bool imageValidated;
    UPGRADE_STATE_T state;
} UPGRADE_DATA_T;

UPGRADE_DATA_T UPGRADE_RUNDATA;

/********************  PRIVATE FUNCTIONS  **************************/
static void sourceUpgradeMessageHandler(Task task, MessageId id, Message message);
static void sourceUpgradeHandleErase(void);

static void sourceUpgradeMessageHandler(Task task, MessageId id, Message message)
{
    UNUSED(task);
    UNUSED(message);
    UNUSED(id);
}

static uint16 ByteUtilsMemCpy16(uint8 *dst, uint16 dstIndex, const uint16 *src, uint16 srcIndex, uint16 size)
{
    uint16 i = 0;

    for(i = 0; i < size; ++i)
    {
        if(i%2)
        {
            dst[i+dstIndex] = src[(i+srcIndex)/2] & 0xFF;
        }
        else
        {
            dst[i+dstIndex] = src[(i+srcIndex)/2] >> 8 ;
        }
    }

    return size;
}

static void sourceUpgradeSendReply(uint8 cmd, bool ack)
{
    uint8 report[USB_DATA_LINK_REPORT_SIZE];

    memset(report, 0, USB_DATA_LINK_REPORT_SIZE);

    report[0] = USB_HOST_COMMAND_FW_UPGRADE;
    report[1] = cmd | (ack?ACK_MASK:0);

    usbHidSendInputReport(USB_DATA_LINK_REPORT_ID, USB_DATA_LINK_REPORT_SIZE, report);
}

static uint32 sourceUpgradeCalcCheckSum(const uint8 *data, uint8 dataSize)
{
    uint32 checkSum = 0;
    uint8 i;

    for(i = 0; i < dataSize; i++)
    {
        checkSum += (uint32)data[i];
    }

    return checkSum;
}

static void sourceUpgradeHandleErase(void)
{
    if(UPGRADE_RUNDATA.state == UPGRADE_STATE_ERASING_FLASH)
    {
        return;
    }
    /* Erase entire area */
    ImageUpgradeErase();

    UPGRADE_RUNDATA.state = UPGRADE_STATE_ERASING_FLASH;
}

static void sourceUpgradeWriteFlash(Sink sink, const uint8 *data, uint16 dataSize)
{
    if(dataSize)
    {
        Source src = StreamRegionSource(data, dataSize);
        do
        {
            /* In case the sink is smaller than the size of incoming data */
            if (dataSize > SinkSlack(sink))
            {
                dataSize = SinkSlack(sink);
                if(!dataSize)
                {
                    SourceClose(src);
                    break;
                }
            }

            if(!StreamMove(sink, src, dataSize))
            {
                break;
            }

            if(!SinkFlushBlocking(sink, dataSize))
            {
                break;
            }

            dataSize = SourceSize(src);

        }while(dataSize);
    }
}

static void sourceUpgradeHandleSelectPartition(UpgradePartitionSelectMsg_t *partitionSelect)
{
    uint32 firstWord;

    if(UPGRADE_RUNDATA.upgradeSink)
    {
        SinkClose(UPGRADE_RUNDATA.upgradeSink);
        UPGRADE_RUNDATA.upgradeSink = NULL;
    }

    UPGRADE_RUNDATA.selectedPartition = partitionSelect->partitionNumber;

    UPGRADE_RUNDATA.selectedpartitionCheckSum  = (uint32)(partitionSelect->checkSum[0]);
    UPGRADE_RUNDATA.selectedpartitionCheckSum |= (uint32)(partitionSelect->checkSum[1] << 8);
    UPGRADE_RUNDATA.selectedpartitionCheckSum |= (uint32)(partitionSelect->checkSum[2] << 16);
    UPGRADE_RUNDATA.selectedpartitionCheckSum |= (uint32)(partitionSelect->checkSum[3] << 24);

    UPGRADE_RUNDATA.selectedpartitionSize  = (uint32)(partitionSelect->partitionSize[0]);
    UPGRADE_RUNDATA.selectedpartitionSize |= (uint32)(partitionSelect->partitionSize[1] << 8);
    UPGRADE_RUNDATA.selectedpartitionSize |= (uint32)(partitionSelect->partitionSize[2] << 16);
    UPGRADE_RUNDATA.selectedpartitionSize |= (uint32)(partitionSelect->partitionSize[3] << 24);

    UPGRADE_RUNDATA.currentPartitionPos = 0;
    UPGRADE_RUNDATA.currentPartitionCheckSum = 0;

    if(UPGRADE_RUNDATA.selectedPartition == PARTITION_SIGNATURE)
    {
        UPGRADE_RUNDATA.signature = (uint8 *)malloc(UPGRADE_RUNDATA.selectedpartitionSize);
        memset(UPGRADE_RUNDATA.signature, 0, UPGRADE_RUNDATA.selectedpartitionSize);
        UPGRADE_RUNDATA.state = UPGRADE_STATE_DOWNLOADING;
        sourceUpgradeSendReply(UPGRADE_SELECT_PARTITION, ACK);
    }
    else
    {
        firstWord  = (uint32)(partitionSelect->firstWord[0]);
        firstWord |= (uint32)(partitionSelect->firstWord[1] << 8);
        firstWord |= (uint32)(partitionSelect->firstWord[2] << 16);
        firstWord |= (uint32)(partitionSelect->firstWord[3] << 24);

        UPGRADE_RUNDATA.upgradeSink = ImageUpgradeStreamGetSink(0, UPGRADE_RUNDATA.selectedPartition, firstWord);

        if(UPGRADE_RUNDATA.upgradeSink)
        {
            UPGRADE_RUNDATA.currentPartitionPos += 4;
            UPGRADE_RUNDATA.currentPartitionCheckSum = sourceUpgradeCalcCheckSum(partitionSelect->firstWord, 4);
            UPGRADE_RUNDATA.state = UPGRADE_STATE_DOWNLOADING;
        }

        sourceUpgradeSendReply(UPGRADE_SELECT_PARTITION, (UPGRADE_RUNDATA.upgradeSink != NULL));
    }
}

static void sourceUpgradeHandleWriteImage(UpgradeImageWriteMsg_t *imageWrite)
{
    if(UPGRADE_RUNDATA.state == UPGRADE_STATE_DOWNLOADING)
    {
        if(UPGRADE_RUNDATA.selectedPartition == PARTITION_SIGNATURE)
        {
            memmove(UPGRADE_RUNDATA.signature + UPGRADE_RUNDATA.currentPartitionPos, imageWrite->data, imageWrite->dataSize);
        }
        else
        {
            sourceUpgradeWriteFlash(UPGRADE_RUNDATA.upgradeSink, imageWrite->data, imageWrite->dataSize);
        }

        UPGRADE_RUNDATA.currentPartitionPos += imageWrite->dataSize;
        UPGRADE_RUNDATA.currentPartitionCheckSum += sourceUpgradeCalcCheckSum(imageWrite->data, imageWrite->dataSize);

        if(imageWrite->lastPacket)
        {
            if(UPGRADE_RUNDATA.currentPartitionPos != UPGRADE_RUNDATA.selectedpartitionSize)
            {
                Panic();
            }

            if(UPGRADE_RUNDATA.currentPartitionCheckSum != UPGRADE_RUNDATA.selectedpartitionCheckSum)
            {
                Panic();
            }

            if(UPGRADE_RUNDATA.selectedPartition == PARTITION_SIGNATURE)
            {
                UPGRADE_RUNDATA.signatureSize = UPGRADE_RUNDATA.selectedpartitionSize;
            }
            else
            {
                SinkClose(UPGRADE_RUNDATA.upgradeSink);
                UPGRADE_RUNDATA.upgradeSink = NULL;
                UPGRADE_RUNDATA.partitionWritten |= (1 << UPGRADE_RUNDATA.selectedPartition);
            }
        }

        sourceUpgradeSendReply(UPGRADE_WRITE_IMAGE, ACK);
    }
    else
    {
        sourceUpgradeSendReply(UPGRADE_WRITE_IMAGE, NACK);
    }
}


static bool sourceUpgradeHandleValidateImagePhase1(void)
{
    hash_context_t vctx;

    if(UPGRADE_RUNDATA.signature && UPGRADE_RUNDATA.signatureSize)
    {
        /* Create the hash context in Apps P0. */
        vctx = ImageUpgradeHashInitialise(SHA256_ALGORITHM);

        if(!vctx)
        {
            Panic();
        }

        ImageUpgradeHashAllSectionsUpdate(vctx);

        UPGRADE_RUNDATA.vctx = vctx;

        return TRUE;
    }

    return FALSE;
}

static bool sourceUpgradeHandleValidateImagePhase2(void)
{
    uint8 sectionHash[SHA_256_HASH_LENGTH];
    uint16 signR2n[RSA_SIGNATURE_SIZE];
    uint16 workSignature[RSA_SIGNATURE_SIZE];
    uint16 reworkSignature[RSA_SIGNATURE_SIZE];
    hash_context_t vctx;
    int verifyResult;

    /* sanity check to make sure the use of defines for array sizes in the function
     * matches the original use of sizeof(). sizeof() removed in case malloc used.
     */
    COMPILE_TIME_ASSERT(sizeof(sectionHash) == SHA_256_HASH_LENGTH,
            compiler_assumptions_changed_u8);
    COMPILE_TIME_ASSERT(sizeof(workSignature) == (sizeof(uint16) * RSA_SIGNATURE_SIZE),
            compiler_assumptions_changed_u16);

    if(UPGRADE_RUNDATA.signature && UPGRADE_RUNDATA.signatureSize)
    {
        vctx = UPGRADE_RUNDATA.vctx;

        /* Get the result of the ImageUpgradeHashSectionUpdate */
        if(!ImageUpgradeHashFinalise(vctx, sectionHash, SHA_256_HASH_LENGTH))
        {
            return FALSE;
        }

        /*
         * Copy the encrypted padded signature given into the workSignature array
         * as the rsa_decrypt will modify the workSignature array to contain the
         * output decrypted PSS padded SHA-256 signature, and we don't want to
         * trample on the original input supplied.
         */
        ByteUtilsMemCpy16((uint8 *)workSignature, 0, (const uint16 *) UPGRADE_RUNDATA.signature, 0,
                            RSA_SIGNATURE_SIZE * sizeof(uint16));

        /*
        * Copy the constant rsa_decrypt_constant_sign_r2n array into a writable
        * array as it will get modified in the rsa_decrypt process.
        */
        memcpy(signR2n, rsa_decrypt_constant_sign_r2n, RSA_SIGNATURE_SIZE * sizeof(uint16));

        /*
        * Decrypt the RSA-2048 encrypted PSS padded signature.
        * The result is the PSS padded signature returned in the signature
        * array that is also used for the input.
        */
        rsa_decrypt(workSignature, &rsa_decrypt_constant_mod, signR2n);

        /*
         * The ce_pkcs1_pss_padding_verify was failing on looking for 0xbc in the
         * last byte of the workSignature output by rsa_decrypt, and it was in the
         * penultimate byte, so swap the uint16 endian-ness from workSignature into
         * reworkSignature and supply that as input to ce_pkcs1_pss_padding_verify
         * instead.
         */
        ByteUtilsMemCpy16((uint8 *)reworkSignature, 0, (const uint16 *) workSignature, 0,
                        RSA_SIGNATURE_SIZE * sizeof(uint16));
        /*
        * Verify the PSS padded signature against the
        * SHA-256 hash in from the image section in the sectionHash.
        */
        verifyResult = ce_pkcs1_pss_padding_verify(
                                  (const unsigned char *)sectionHash,
                                  SHA_256_HASH_LENGTH,
                                  (const unsigned char *)reworkSignature,
                                  RSA_SIGNATURE_SIZE * sizeof(uint16),
                                  ce_pkcs1_pss_padding_verify_constant_saltlen,
                                  sizeof(rsa_decrypt_constant_mod.M) * 8);

        if(verifyResult != CE_SUCCESS)
        {
            return FALSE;
        }

        UPGRADE_RUNDATA.imageValidated = TRUE;

        return TRUE;
    }

    return FALSE;
}

static void sourceUpgradeHandleReadVersion(void)
{
    uint8 versionMajor = 0, versionMinor = 0, versionPatch = 0;
    uint8 report[USB_DATA_LINK_REPORT_SIZE];

    sourcePsGetVersion(&versionMajor, &versionMinor, &versionPatch);

    memset(report, 0, USB_DATA_LINK_REPORT_SIZE);

    report[0] = USB_HOST_COMMAND_FW_UPGRADE;
    report[1] = UPGRADE_READ_VERSION | ACK_MASK;
    report[2] = versionMajor;
    report[3] = versionMinor;
    report[4] = versionPatch;

    usbHidSendInputReport(USB_DATA_LINK_REPORT_ID, USB_DATA_LINK_REPORT_SIZE, report);
}

static void sourceUpgradeHandleWriteVersion(uint8 versionMajor, uint8 versionMinor, uint8 versionPatch)
{
    UNUSED(versionMajor);
    UNUSED(versionMinor);
    UNUSED(versionPatch);

    sourceUpgradeSendReply(UPGRADE_WRITE_VERSION, ACK);
}

static void sourceUpgradeHandleCopyMissingSections(void)
{
    if(UPGRADE_RUNDATA.imageValidated)
    {
        ImageUpgradeCopy();
    }
    else
    {
        sourceUpgradeSendReply(UPGRADE_COPY_MISSING_SECTIONS, FALSE);
    }
}

static void sourceUpgradeCommit(void)
{
    sourcePsSetFwUpgrade(FALSE);

    if(!ImageUpgradeSwapTryStatus())
    {
        BootSetMode(BootGetMode());
        return;
    }

    if(!ImageUpgradeSwapCommit())
    {
        BootSetMode(BootGetMode());
        return;
    }
}

void sourceUpgradeInit(void)
{
    UPGRADE_RUNDATA.upgradeTask.handler = sourceUpgradeMessageHandler;
    UPGRADE_RUNDATA.state = UPGRADE_STATE_IDLE;
    UPGRADE_RUNDATA.selectedPartition = IMAGE_SECTION_ID_MAX;
    UPGRADE_RUNDATA.partitionWritten = 0;
    UPGRADE_RUNDATA.imageValidated = FALSE;

    if(sourcePsGetFwUpgrade())
    {
        sourceUpgradeCommit();
    }
}

void sourceUpgradeHandleHidMessage(uint8* data, uint8 data_size)
{
    UNUSED(data_size);
    uint8 command = *data;

    switch(command)
    {
    case UPGRADE_ERASE_FLASH:
        sourceUpgradeHandleErase();
        break;
    case UPGRADE_SELECT_PARTITION:
        sourceUpgradeHandleSelectPartition((UpgradePartitionSelectMsg_t*)&data[1]);
        break;
    case UPGRADE_WRITE_IMAGE:
        sourceUpgradeHandleWriteImage((UpgradeImageWriteMsg_t *)&data[1]);
        break;
    case UPGRADE_VALIDATE_IMAGE:
        if(!sourceUpgradeHandleValidateImagePhase1())
        {
            sourceUpgradeSendReply(UPGRADE_VALIDATE_IMAGE, FALSE);
        }
        break;
    case UPGRADE_READ_VERSION:
        sourceUpgradeHandleReadVersion();
        break;
    case UPGRADE_WRITE_VERSION:
        sourceUpgradeHandleWriteVersion(data[1], data[2], data[3]);
        break;
    case UPGRADE_COPY_MISSING_SECTIONS:
        sourceUpgradeHandleCopyMissingSections();
        break;
    default:
        break;
    }
}

bool appUpgradeHandleSystemMessages(MessageId id, Message message, bool already_handled)
{
    switch (id)
    {
        case MESSAGE_IMAGE_UPGRADE_ERASE_STATUS:
        {
            bool status = ((MessageImageUpgradeEraseStatus *)message)->erase_status;
            if(status)
            {
                /* erase success */
                sourceUpgradeSendReply(UPGRADE_ERASE_FLASH, ACK);
            }
        }
        break;

        case MESSAGE_IMAGE_UPGRADE_COPY_STATUS:
        {
            /* Missing sections in an image update are now copied from
              the boot image bank to the new image bank */
            bool status = ((MessageImageUpgradeCopyStatus *)message)->copy_status;
            if(status)
            {
                 /* Let application know that copy is done */
                sourceUpgradeSendReply(UPGRADE_COPY_MISSING_SECTIONS, TRUE);
                /* Set fw upgrade flag as there are some things need to do after reboot */
                sourcePsSetFwUpgrade(TRUE);
                /* Reboot and use new image bank */
                ImageUpgradeSwapTry();
            }
        }
        break;

        case MESSAGE_IMAGE_UPGRADE_AUDIO_STATUS:
        {
        }
        break;

        case MESSAGE_IMAGE_UPGRADE_HASH_ALL_SECTIONS_UPDATE_STATUS:
        {
            MessageImageUpgradeHashAllSectionsUpdateStatus *msg = (MessageImageUpgradeHashAllSectionsUpdateStatus*)message;

            if(msg->status && sourceUpgradeHandleValidateImagePhase2())
            {
                sourceUpgradeSendReply(UPGRADE_VALIDATE_IMAGE, TRUE);
            }
            else
            {
                sourceUpgradeSendReply(UPGRADE_VALIDATE_IMAGE, FALSE);
            }
        }
        break;
    }

    return already_handled;
}

#endif /* INCLUDE_DFU */
